#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Betaflight PID参数分析和优化工具
基于飞行记录数据分析当前PID性能并提供优化建议
"""

import os
import re
import numpy as np
import matplotlib.pyplot as plt
from orangebox import Parser
import matplotlib.font_manager as fm

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False

class PIDOptimizer:
    def __init__(self, config_file, log_files):
        self.config_file = config_file
        self.log_files = log_files
        self.current_pids = {}
        self.flight_data = {}
        self.analysis_results = {}
        
    def parse_config(self):
        """解析当前飞控配置文件"""
        print("解析飞控配置文件...")
        
        with open(self.config_file, 'r', encoding='utf-8') as f:
            content = f.read()
            
        # 提取PID参数
        pid_patterns = {
            'p_roll': r'set p_roll = (\d+)',
            'i_roll': r'set i_roll = (\d+)',
            'd_roll': r'set d_roll = (\d+)',
            'f_roll': r'set f_roll = (\d+)',
            'p_pitch': r'set p_pitch = (\d+)',
            'i_pitch': r'set i_pitch = (\d+)',
            'd_pitch': r'set d_pitch = (\d+)',
            'f_pitch': r'set f_pitch = (\d+)',
            'p_yaw': r'set p_yaw = (\d+)',
            'i_yaw': r'set i_yaw = (\d+)',
            'd_yaw': r'set d_yaw = (\d+)',
            'f_yaw': r'set f_yaw = (\d+)',
        }
        
        for param, pattern in pid_patterns.items():
            match = re.search(pattern, content)
            if match:
                self.current_pids[param] = int(match.group(1))
                
        # 提取其他重要参数
        other_patterns = {
            'anti_gravity_gain': r'set anti_gravity_gain = (\d+)',
            'iterm_relax_cutoff': r'set iterm_relax_cutoff = (\d+)',
            'dterm_lpf1_static_hz': r'set dterm_lpf1_static_hz = (\d+)',
            'dterm_lpf2_static_hz': r'set dterm_lpf2_static_hz = (\d+)',
            'gyro_lpf1_static_hz': r'set gyro_lpf1_static_hz = (\d+)',
            'gyro_lpf2_static_hz': r'set gyro_lpf2_static_hz = (\d+)',
        }
        
        for param, pattern in other_patterns.items():
            match = re.search(pattern, content)
            if match:
                self.current_pids[param] = int(match.group(1))
                
        print(f"当前PID配置: {self.current_pids}")
        
    def analyze_flight_data(self):
        """分析飞行记录数据"""
        print("分析飞行记录数据...")
        
        for log_file in self.log_files:
            if not os.path.exists(log_file):
                continue
                
            print(f"分析文件: {log_file}")
            
            try:
                parser = Parser.load(log_file)
                data = {}
                field_names = parser.field_names
                
                # 初始化数据字典
                for field in field_names:
                    data[field] = []
                
                # 提取数据
                for frame in parser.frames():
                    if frame.data is not None:
                        for i, field in enumerate(field_names):
                            if i < len(frame.data) and frame.data[i] is not None:
                                data[field].append(frame.data[i])
                            else:
                                data[field].append(0)
                
                # 转换为numpy数组
                for field in data:
                    data[field] = np.array(data[field])
                
                self.flight_data[os.path.basename(log_file)] = data
                
            except Exception as e:
                print(f"分析文件 {log_file} 时出错: {e}")
                
    def calculate_performance_metrics(self):
        """计算性能指标"""
        print("计算性能指标...")
        
        for filename, data in self.flight_data.items():
            metrics = {}
            
            # 检查必要字段是否存在
            required_fields = ['setpoint[0]', 'setpoint[1]', 'setpoint[2]', 
                             'gyroADC[0]', 'gyroADC[1]', 'gyroADC[2]']
            
            available_fields = list(data.keys())
            
            # 寻找相似字段
            setpoint_fields = [f for f in available_fields if 'setpoint' in f.lower()]
            gyro_fields = [f for f in available_fields if 'gyro' in f.lower()]
            
            if len(setpoint_fields) >= 3 and len(gyro_fields) >= 3:
                # Roll轴分析
                if len(setpoint_fields) > 0 and len(gyro_fields) > 0:
                    setpoint_roll = data[setpoint_fields[0]]
                    actual_roll = data[gyro_fields[0]]
                    
                    error_roll = setpoint_roll - actual_roll
                    metrics['roll_rms_error'] = np.sqrt(np.mean(error_roll**2))
                    metrics['roll_max_error'] = np.max(np.abs(error_roll))
                    metrics['roll_settling_time'] = self.calculate_settling_time(error_roll)
                    
                # Pitch轴分析
                if len(setpoint_fields) > 1 and len(gyro_fields) > 1:
                    setpoint_pitch = data[setpoint_fields[1]]
                    actual_pitch = data[gyro_fields[1]]
                    
                    error_pitch = setpoint_pitch - actual_pitch
                    metrics['pitch_rms_error'] = np.sqrt(np.mean(error_pitch**2))
                    metrics['pitch_max_error'] = np.max(np.abs(error_pitch))
                    metrics['pitch_settling_time'] = self.calculate_settling_time(error_pitch)
                    
                # Yaw轴分析
                if len(setpoint_fields) > 2 and len(gyro_fields) > 2:
                    setpoint_yaw = data[setpoint_fields[2]]
                    actual_yaw = data[gyro_fields[2]]
                    
                    error_yaw = setpoint_yaw - actual_yaw
                    metrics['yaw_rms_error'] = np.sqrt(np.mean(error_yaw**2))
                    metrics['yaw_max_error'] = np.max(np.abs(error_yaw))
                    metrics['yaw_settling_time'] = self.calculate_settling_time(error_yaw)
            
            # 振荡分析
            for axis, fields in [('roll', gyro_fields[:1]), ('pitch', gyro_fields[1:2]), ('yaw', gyro_fields[2:3])]:
                if fields:
                    gyro_data = data[fields[0]]
                    metrics[f'{axis}_oscillation'] = self.detect_oscillation(gyro_data)
            
            self.analysis_results[filename] = metrics
            
    def calculate_settling_time(self, error_signal, threshold=0.05):
        """计算稳定时间"""
        if len(error_signal) == 0:
            return 0
            
        # 找到误差信号稳定在阈值内的时间
        stable_mask = np.abs(error_signal) < threshold * np.max(np.abs(error_signal))
        
        if np.any(stable_mask):
            # 找到最后一次超出阈值的位置
            last_unstable = np.where(~stable_mask)[0]
            if len(last_unstable) > 0:
                return len(error_signal) - last_unstable[-1]
        
        return len(error_signal)
    
    def detect_oscillation(self, signal, window_size=50):
        """检测振荡"""
        if len(signal) < window_size:
            return 0
            
        # 计算信号的标准差作为振荡指标
        oscillation_values = []
        for i in range(len(signal) - window_size):
            window = signal[i:i+window_size]
            oscillation_values.append(np.std(window))
            
        return np.mean(oscillation_values) if oscillation_values else 0
    
    def generate_optimization_recommendations(self):
        """生成优化建议"""
        print("生成优化建议...")
        
        recommendations = []
        recommendations.append("# Betaflight PID优化建议")
        recommendations.append(f"# 基于飞行记录分析生成 - {self.__class__.__name__}")
        recommendations.append("")
        
        # 当前配置
        recommendations.append("## 当前PID配置")
        recommendations.append(f"Roll轴: P={self.current_pids.get('p_roll', 0)}, I={self.current_pids.get('i_roll', 0)}, D={self.current_pids.get('d_roll', 0)}, F={self.current_pids.get('f_roll', 0)}")
        recommendations.append(f"Pitch轴: P={self.current_pids.get('p_pitch', 0)}, I={self.current_pids.get('i_pitch', 0)}, D={self.current_pids.get('d_pitch', 0)}, F={self.current_pids.get('f_pitch', 0)}")
        recommendations.append(f"Yaw轴: P={self.current_pids.get('p_yaw', 0)}, I={self.current_pids.get('i_yaw', 0)}, D={self.current_pids.get('d_yaw', 0)}, F={self.current_pids.get('f_yaw', 0)}")
        recommendations.append("")
        
        # 性能分析
        recommendations.append("## 性能分析结果")
        for filename, metrics in self.analysis_results.items():
            recommendations.append(f"### {filename}")
            for metric, value in metrics.items():
                recommendations.append(f"{metric}: {value:.4f}")
            recommendations.append("")
        
        # 优化建议
        recommendations.append("## 优化建议")
        
        # 基于分析结果提供建议
        avg_metrics = self.calculate_average_metrics()
        
        if avg_metrics:
            # Roll轴建议
            if avg_metrics.get('roll_oscillation', 0) > 10:
                recommendations.append("### Roll轴优化")
                recommendations.append("- 检测到Roll轴振荡，建议降低D项参数")
                new_d_roll = max(20, int(self.current_pids.get('d_roll', 35) * 0.8))
                recommendations.append(f"- 建议 d_roll 从 {self.current_pids.get('d_roll', 35)} 调整为 {new_d_roll}")
                recommendations.append("")
            
            if avg_metrics.get('roll_rms_error', 0) > 5:
                recommendations.append("### Roll轴跟踪性能")
                recommendations.append("- Roll轴跟踪误差较大，建议增加P项")
                new_p_roll = min(80, int(self.current_pids.get('p_roll', 48) * 1.1))
                recommendations.append(f"- 建议 p_roll 从 {self.current_pids.get('p_roll', 48)} 调整为 {new_p_roll}")
                recommendations.append("")
            
            # Pitch轴建议
            if avg_metrics.get('pitch_oscillation', 0) > 10:
                recommendations.append("### Pitch轴优化")
                recommendations.append("- 检测到Pitch轴振荡，建议降低D项参数")
                new_d_pitch = max(20, int(self.current_pids.get('d_pitch', 32) * 0.8))
                recommendations.append(f"- 建议 d_pitch 从 {self.current_pids.get('d_pitch', 32)} 调整为 {new_d_pitch}")
                recommendations.append("")
            
            if avg_metrics.get('pitch_rms_error', 0) > 5:
                recommendations.append("### Pitch轴跟踪性能")
                recommendations.append("- Pitch轴跟踪误差较大，建议增加P项")
                new_p_pitch = min(80, int(self.current_pids.get('p_pitch', 52) * 1.1))
                recommendations.append(f"- 建议 p_pitch 从 {self.current_pids.get('p_pitch', 52)} 调整为 {new_p_pitch}")
                recommendations.append("")
            
            # Yaw轴建议
            if avg_metrics.get('yaw_oscillation', 0) > 8:
                recommendations.append("### Yaw轴优化")
                recommendations.append("- 检测到Yaw轴振荡，建议降低P项参数")
                new_p_yaw = max(40, int(self.current_pids.get('p_yaw', 65) * 0.9))
                recommendations.append(f"- 建议 p_yaw 从 {self.current_pids.get('p_yaw', 65)} 调整为 {new_p_yaw}")
                recommendations.append("")
        
        # 通用建议
        recommendations.append("### 通用建议")
        recommendations.append("- 建议逐步调整参数，每次只调整一个轴的一个参数")
        recommendations.append("- 调整后进行试飞验证，观察飞行表现")
        recommendations.append("- 如果出现振荡，优先降低D项参数")
        recommendations.append("- 如果跟踪性能不足，可适当增加P项参数")
        recommendations.append("- I项参数通常不需要大幅调整")
        recommendations.append("")
        
        return recommendations
    
    def calculate_average_metrics(self):
        """计算平均性能指标"""
        if not self.analysis_results:
            return {}
            
        all_metrics = {}
        for metrics in self.analysis_results.values():
            for key, value in metrics.items():
                if key not in all_metrics:
                    all_metrics[key] = []
                all_metrics[key].append(value)
        
        avg_metrics = {}
        for key, values in all_metrics.items():
            avg_metrics[key] = np.mean(values)
            
        return avg_metrics
    
    def generate_optimized_config(self):
        """生成优化后的配置文件"""
        print("生成优化配置...")
        
        # 基于分析结果计算优化参数
        optimized_pids = self.current_pids.copy()
        avg_metrics = self.calculate_average_metrics()
        
        if avg_metrics:
            # Roll轴优化
            if avg_metrics.get('roll_oscillation', 0) > 10:
                optimized_pids['d_roll'] = max(20, int(optimized_pids.get('d_roll', 35) * 0.8))
            if avg_metrics.get('roll_rms_error', 0) > 5:
                optimized_pids['p_roll'] = min(80, int(optimized_pids.get('p_roll', 48) * 1.1))
            
            # Pitch轴优化
            if avg_metrics.get('pitch_oscillation', 0) > 10:
                optimized_pids['d_pitch'] = max(20, int(optimized_pids.get('d_pitch', 32) * 0.8))
            if avg_metrics.get('pitch_rms_error', 0) > 5:
                optimized_pids['p_pitch'] = min(80, int(optimized_pids.get('p_pitch', 52) * 1.1))
            
            # Yaw轴优化
            if avg_metrics.get('yaw_oscillation', 0) > 8:
                optimized_pids['p_yaw'] = max(40, int(optimized_pids.get('p_yaw', 65) * 0.9))
        
        # 生成CLI命令
        config_commands = []
        config_commands.append("# 优化后的PID参数")
        config_commands.append("# 请在Betaflight Configurator的CLI中执行以下命令")
        config_commands.append("")
        
        for param, value in optimized_pids.items():
            if param in ['p_roll', 'i_roll', 'd_roll', 'f_roll', 'p_pitch', 'i_pitch', 'd_pitch', 'f_pitch', 'p_yaw', 'i_yaw', 'd_yaw', 'f_yaw']:
                config_commands.append(f"set {param} = {value}")
        
        config_commands.append("")
        config_commands.append("save")
        
        return config_commands
    
    def run_analysis(self):
        """运行完整分析"""
        print("开始PID优化分析...")
        
        # 解析配置
        self.parse_config()
        
        # 分析飞行数据
        self.analyze_flight_data()
        
        # 计算性能指标
        self.calculate_performance_metrics()
        
        # 生成建议
        recommendations = self.generate_optimization_recommendations()
        
        # 生成优化配置
        optimized_config = self.generate_optimized_config()
        
        return recommendations, optimized_config

def main():
    # 配置文件路径
    config_file = "d:/Git/FlightRecord/SFPV_85/BTFL_BETAFPVF405_飞控当前配置.txt"
    
    # 日志文件路径
    log_files = [
        "d:/Git/FlightRecord/SFPV_85/record/btfl_001.bbl",
        "d:/Git/FlightRecord/SFPV_85/record/btfl_002.bbl",
        "d:/Git/FlightRecord/SFPV_85/record/btfl_all.bbl"
    ]
    
    # 创建优化器
    optimizer = PIDOptimizer(config_file, log_files)
    
    # 运行分析
    recommendations, optimized_config = optimizer.run_analysis()
    
    # 保存结果
    output_dir = "d:/Git/FlightRecord/SFPV_85/optimization_results"
    os.makedirs(output_dir, exist_ok=True)
    
    # 保存建议
    with open(os.path.join(output_dir, "PID优化建议.md"), 'w', encoding='utf-8') as f:
        f.write('\n'.join(recommendations))
    
    # 保存优化配置
    with open(os.path.join(output_dir, "优化后的PID配置.txt"), 'w', encoding='utf-8') as f:
        f.write('\n'.join(optimized_config))
    
    print(f"\n分析完成！结果已保存到: {output_dir}")
    print("\n=== PID优化建议 ===")
    for line in recommendations[-20:]:  # 显示最后20行建议
        print(line)

if __name__ == "__main__":
    main()